import torch
import mindspore as ms

# /root/linx/segan_pt/pt_disciminator.pth
pt_path = "/root/linx/segan_pt/pt_generator.pth"
# /root/linx/segan_ms/ms_disciminator.ckpt
ms_path = "/root/linx/segan_ms/ms_generator.ckpt"

pt_ckpt = torch.load(pt_path, map_location='cpu')
ms_ckpt = ms.load_checkpoint(ms_path)

new_param_list = []
for name, param in pt_ckpt.items():
    new_name = 'G.' + name
    if "conv.weight" in new_name:
        param = param.unsqueeze(2)
    if "act.weight" in new_name:
        new_name = new_name.replace("act.weight", "act.w")
    new_param = ms.Tensor(param.numpy())
    new_param_list.append({"name": new_name, "data": new_param})

ms.save_checkpoint(new_param_list, "converted_g.ckpt")
pass
